import numpy as np
from scipy import stats
import matplotlib.pyplot as plt

def sampling_and_empirical_dis():
    xk = np.arange(7)  # 所有可能的取值
    print(xk)  # [0 1 2 3 4 5 6]
    pk = (0.1, 0.2, 0.3, 0.1, 0.1, 0.0, 0.2)  # 各个取值的概率
    custm = stats.rv_discrete(name='custm', values=(xk, pk))

    X1 = custm.rvs(size=20)  # 第一次抽样
    X2 = custm.rvs(size=200)  # 第二次抽样
    # 计算X1＆X2中各个结果出现的频率(相当于PMF)
    val1, cnt1 = np.unique(X1, return_counts=True)
    val2, cnt2 = np.unique(X2, return_counts=True)
    pmf_X1 = cnt1 / len(X1)
    pmf_X2 = cnt2 / len(X2)

    plt.figure(1)
    plt.subplot(211)
    plt.plot(xk, custm.pmf(xk), 'ro', ms=8, mec='r', label='theor. pmf')
    plt.vlines(xk, 0, custm.pmf(xk), colors='r', lw=5, alpha=0.2)
    plt.vlines(val1, 0, pmf_X1, colors='b', linestyles='-', lw=3, label='X1 empir. pmf')
    plt.legend(loc='best', frameon=False)
    plt.ylabel('Probability')
    plt.title('Theoretical dist. PMF vs Empirical dist. PMF')
    plt.subplot(212)
    plt.plot(xk, custm.pmf(xk), 'ro', ms=8, mec='r', label='theor. pmf')
    plt.vlines(xk, 0, custm.pmf(xk), colors='r', lw=5, alpha=0.2)
    plt.vlines(val2, 0, pmf_X2, colors='g', linestyles='-', lw=3, label='X2 empir. pmf')
    plt.legend(loc='best', frameon=False)
    plt.ylabel('Probability')
    plt.show()

sampling_and_empirical_dis()